# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: yuxuan

import datetime
import functools
import logging
import signal

from absl import (flags, app)
from tornado.ioloop import IOLoop, PeriodicCallback
import numpy as np

import coin.flow.subscriber as flow_subscriber
from coin.exchange.kr_rest.product_holders import get_holder_from_product
from coin.exchange.base.order_gateway import OrderSide, OrderType, Order
from coin.strategy.hk_hard_arb.products import NativeProducts
from coin.exchange.base.kr_rest.product import NativeProduct
from coin.exchange.base.kr_rest.futures_product import NativeFuturesProduct
from coin.exchange.base.strategy_logger import StrategyLogProvider
from coin.base.book.book_builder_base import BookBuilderBase
from coin.base.param_util import to_list
from coin.exchange.shared.kr_rest.product import SharedProduct
from coin.exchange.util.feed_checker import get_default_feed_checker_config
from coin.flow.simple_queue import SimpleQueue
from coin.proto.coin_order_checker_pb2 import AccountCheckerConfig
from coin.strategy.executor.util.cancel_all_executor import (SimpleCancelAllExecutor)
from coin.strategy.accounting.pnl_balance_printer.factory import get_pnl_balance_printer_cls
from coin.strategy.accounting.pnl_balance_printer.util import reserve_as_proto
from coin.strategy.mm.base.bookmap import BookMap
from coin.strategy.hk_hard_arb.util import load_clean_json
from coin.strategy.mm.subscription import FeedSubscriptionRequest
from coin.strategy.mm.feed import SubscriptionRequest, get_flow_book_builder_cls
from coin.strategy.mm.live_envs import (create_live_env_instance, get_live_env_cls)
from coin.strategy.executor.aggressive.factory import construct_agg_executor
from coin.strategy.executor.aggressive.util import \
  convert_executor_config_from_dict as convert_agg_config
from coin.support.accounting.checker.account_checker import AccountChecker
from coin.support.accounting.checker.config_util import convert_checker_config_from_dict
from coin.exchange.base.order_gateway_config import convert_og_config_from_dict
from coin.exchange.huobi_futures.kr_rest.native_private_client import HuobiFuturesOrderSide
from coin.strategy.tool.fill_log_dumper import FillLogDumper
from coin.strategy.basis_strat.fill_info_writer import FillInfo, FillInfoWriter
from coin.strategy.util.request_util import convert_strategy_request_from_dict
from coin.strategy.accounting.fee import get_fee_rate
from coin.strategy.basis_strat.moving_average import MovingAverage
from coin.strategy.basis_strat.base_strategy import BaseStrategy

FLAGS = flags.FLAGS


class DummyPosition(object):
  def __init__(self):
    self.long = self.short = 0


def _gen_arb_products(common, products0, products1, ProductType0, ProductType1):
  if 'exch0_quote_currency' in common and 'exch1_quote_currency' in common:
    exch0_quote_currency = common['exch0_quote_currency']
    exch1_quote_currency = common['exch1_quote_currency']
  else:
    raise ValueError(common)

  products0 = NativeProducts.FromTypedStrProduct(exch0_quote_currency,
                                                 products0.keys(),
                                                 ProductType0)
  products1 = NativeProducts.FromTypedStrProduct(exch1_quote_currency,
                                                 products1.keys(),
                                                 ProductType1)
  assert len(products0.products) == len(products1.products)
  return products0, products1


class BasisUnhedgedSubStrategy(object):
  def __init__(self, parent, product0, product1, fillinfo_writer, strategy_logger):
    self._product0 = product0
    self._product1 = product1
    self._fillinfo_writer = fillinfo_writer
    self._ticker0 = product0.full_symbol
    self._ticker1 = product1.full_symbol
    self._tick_info1 = (get_holder_from_product(product1).product_info.gen_tick(),
                        get_holder_from_product(product1).product_info.gen_tick_qty())
    if hasattr(self._tick_info1[0], 'tick_size'):
      ticksize = self._tick_info1[0].tick_size
      self._next_price_func = lambda p: p + ticksize - (p % ticksize)
      self._prev_price_func = lambda p: p - (p % ticksize)
    else:
      self._next_price_func = self._tick_info1[0].get_next_price
      self._prev_price_func = self._tick_info1[0].get_prev_price

    self._config1 = None
    self._config1_proto = None
    self.og1 = None
    self._agg_executor1 = None
    self._parent = parent
    self.logger = parent.logger
    self._strategy_logger = strategy_logger
    self.feed = parent._last_feed

    # Params
    self._enter_position_threshold = FLAGS.enter_position_threshold
    self._exit_position_threshold = FLAGS.exit_position_threshold
    self._window_period = FLAGS.window_period_sec * 10**9
    self._ref_price_window_period = FLAGS.ref_price_window_period_sec * 10**9
    self._trading_threshold = FLAGS.trading_threshold_bp * 1e-4
    self._max_std_ratio = FLAGS.max_std_ratio
    self._min_std_ratio = FLAGS.min_std_ratio
    self._spread_multiplier = FLAGS.spread_multiplier

    self._ma = MovingAverage(self._window_period)
    self._mean = None
    self._std = None
    self._book_ts = [0, 0]
    self._book_invalid_period = FLAGS.book_invalid_sec * 10**9
    self._aggressive_amt = FLAGS.aggressive_bp * 1e-4

    self._pos1 = 0
    self._pos1_ts = 0
    self._og_pos1 = 0
    self._og_pos1_ts = 0
    self._update_agg_pos_with_balance = FLAGS.update_agg_pos_with_balance
    self._update_position_duration = FLAGS.update_pos_with_balance_sec * 10**9
    self._update_working_qty_duration = FLAGS.update_working_qty_sec * 10**9
    self._equal_theshold = None

    self._lot_size = None
    self._max_pos = None
    self._min_pos = None
    self._neutral_pos = None
    self._agg_taker_fee = None
    self._expected_working_qty = 0
    self._last_agg_submit_fill_ts = 0
    # Handle agg working order update delay.
    self._agg_working_order_fills = dict()

    # Decide if book has changed.
    self._book_levels = FLAGS.book_levels
    self._last_book0 = None
    self._last_book1 = None

    self._get_pos_func = None
    self._is_buy_order_func = None
    self._dummy_position = DummyPosition()

    self._ref_price_dq = MovingAverage(self._ref_price_window_period)
    self._ref_ma = np.nan

  def get_states(self):
    return {
        "mean": self._mean,
        "std": self._std,
        "current_mean": self._ma.get_latest_data() if self._ma else None,
        "dq_size": self._ma.size() if self._ma else 0
    }

  def _get_deque_file_name(self):
    return '%s_%s_%s' % \
        (self._product0.exchange, self._product1.exchange, self._ticker0)

  def _init_get_pos_func(self, og):
    product_type = self._product1.product_type
    exchange = self._product1.exchange
    if product_type == 'Futures':
      if exchange == 'Huobi':
        self._get_pos_func = lambda: \
            (og.get_huobi_futures_position(self._product1) or self._dummy_position).long - \
            (og.get_huobi_futures_position(self._product1) or self._dummy_position).short
        self._is_buy_order_func = lambda order_side: order_side in [
            HuobiFuturesOrderSide.BUY_OPEN, HuobiFuturesOrderSide.BUY_CLOSE, OrderSide.BUY
        ]
    elif product_type == 'CurrencyPair':
      self._get_pos_func = lambda: og.get_balance(self._product1.base.currency)
      self._is_buy_order_func = lambda order_side: order_side == OrderSide.BUY
    assert self._get_pos_func is not None, '%s %s' % (product_type, exchange)

  def write_executor_config(self):
    if self._config1_proto is not None:
      self._strategy_logger.write_executor_config(self._config1_proto)

  def init_executors(self, req1, config1, og1):
    self.og1 = og1
    self._config1 = config1.copy()
    self._config1['force_fire'] = True
    self._config1_proto = convert_agg_config(self._product1.symbol, self._config1)
    self._agg_executor1 = construct_agg_executor(req1, self._config1_proto, self.og1, self.logger)
    self._init_get_pos_func(og1)
    self._equal_theshold = max(self._config1['lot_size'] * 1e-4, 1e-5)
    self._lot_size = self._config1['lot_size']
    self._max_pos = self._config1['max_pos']
    self._min_pos = self._config1['min_pos']
    self._max_std_ratio = self._config1.get('max_std_ratio', self._max_std_ratio)
    self._min_std_ratio = self._config1.get('min_std_ratio', self._min_std_ratio)
    assert 0 < self._min_std_ratio < 1e-2, self._min_std_ratio
    assert 0 < self._max_std_ratio < 1e-2, self._max_std_ratio
    self._enter_position_threshold = self._config1.get('enter_position_threshold',
                                                       self._enter_position_threshold)
    self._exit_position_threshold = self._config1.get('exit_position_threshold',
                                                      self._exit_position_threshold)
    self._neutral_pos = (self._max_pos + self._min_pos) / 2.0
    _, self._agg_taker_fee = get_fee_rate(self.og1.market_type, self.og1.exchange)
    if self._product1.expiry:
      self.logger.info('Futures product, set fees to 0, %s' % self._product1.expiry)
      self._agg_taker_fee = 0.0
    self._update_pos1(self._get_pos_func())
    self._og_pos1 = self._pos1
    self._og_pos1_ts = self._get_ts()
    assert abs(self._agg_taker_fee) <= 1e-2, self._agg_taker_fee
    self.logger.info('%s init pos1 %f' % (self._ticker0, self._pos1))

  def _update_pos1(self, pos):
    self._pos1 = pos
    self._pos1_ts = self._get_ts()

  def _get_ts(self):
    return self._parent._ts

  def _is_book_valid(self):
    return self._get_ts() - self._book_ts[0] < self._book_invalid_period

  def _update_deque(self, implied_rate, mid0, mid1):
    ts = self._get_ts()
    self._ma.update(ts, implied_rate)
    self._ref_price_dq.update(ts, mid0)

  def get_implied_rate(self):
    return self._ma.get_latest_data()

  def _update_statistics(self):
    self._mean, self._std = self._ma.get_statistics()
    self._ref_ma, _ = self._ref_price_dq.get_statistics()

  def ready(self):
    return self._ma.ready()

  def _position_equal(self, p1, p2):
    return abs(p1 - p2) < self._equal_theshold

  def _get_agg_working_qty(self):
    # Use working orders to update |self._expected_working_qty|.
    if self._get_ts() - self._last_agg_submit_fill_ts > self._update_working_qty_duration:
      order_id_set = set()
      working_qty = 0.0
      for order in self._agg_executor1._wo.buy_working_orders:
        working_qty += (order.qty - self._agg_working_order_fills.get(order.order_id, 0.0))
        order_id_set.add(order.order_id)
      for order in self._agg_executor1._wo.sell_working_orders:
        working_qty -= (order.qty - self._agg_working_order_fills.get(order.order_id, 0.0))
        order_id_set.add(order.order_id)
      # Remove order ids that are not in working orders.
      obsolete_ids = self._agg_working_order_fills.keys() - order_id_set
      for order_id in obsolete_ids:
        del self._agg_working_order_fills[order_id]
      if not self._position_equal(self._expected_working_qty, working_qty):
        self.logger.info('%s working qty expected %f, updated %f, %d %d' %
                         (self._ticker1,
                          self._expected_working_qty,
                          working_qty,
                          len(self._agg_executor1._wo.buy_working_orders),
                          len(self._agg_executor1._wo.sell_working_orders)))
      self._expected_working_qty = working_qty
    return self._expected_working_qty

  def _is_agg_pos_valid(self):
    ts = self._get_ts()
    return ts - max(self._pos1_ts, self._og_pos1_ts) > self._update_position_duration

  def _get_agg_balance(self):
    og_pos = self._get_pos_func()
    if self._update_agg_pos_with_balance:
      if not self._position_equal(og_pos, self._og_pos1):
        self._og_pos1 = og_pos
        self._og_pos1_ts = self._get_ts()
      if not self._position_equal(self._pos1, self._og_pos1) \
          and self._is_agg_pos_valid():
        self.logger.info('Update pos1 with og_pos1, %f %f' % (self._pos1, self._og_pos1))
        self._pos1 = self._og_pos1
    else:
      if self._position_equal(self._pos1, og_pos):
        self._pos1 = og_pos
    return self._pos1

  def _get_target_position_info(self, implied_rate, spread):
    self._update_statistics()
    mean = self._mean
    std = min(self._std, mean * self._max_std_ratio)
    std = max(self._std, mean * self._min_std_ratio)
    cnt_std = (implied_rate - mean) / std
    agg_pos = self._get_agg_balance()
    working_qty = self._get_agg_working_qty()
    potential_agg_pos = agg_pos + working_qty
    target_pos1 = potential_agg_pos
    extra_std = 0.0

    if cnt_std > -self._exit_position_threshold:
      target_pos1 = min(self._neutral_pos, potential_agg_pos)
      extra_std = cnt_std + self._exit_position_threshold
    elif cnt_std < self._exit_position_threshold:
      target_pos1 = max(self._neutral_pos, potential_agg_pos)
      extra_std = self._exit_position_threshold - cnt_std

    opportunity_spread = abs(implied_rate - mean) / mean
    if opportunity_spread > self._spread_multiplier * spread:
      if cnt_std > self._enter_position_threshold:
        target_pos1 = self._min_pos
        extra_std = cnt_std - self._enter_position_threshold
      elif cnt_std < -self._enter_position_threshold:
        target_pos1 = self._max_pos
        extra_std = -self._enter_position_threshold - cnt_std
    return target_pos1 - potential_agg_pos, agg_pos, working_qty, extra_std * std

  def _get_agg_buy_price(self, price, aggressive_amt):
    price = price * (1 + aggressive_amt)
    return self._prev_price_func(price)

  def _get_agg_sell_price(self, price, aggressive_amt):
    price = price * (1 - aggressive_amt)
    return self._next_price_func(price)

  def _process(self, implied_rate, ref_ma, bids0, asks0, bids1, asks1, mid0, mid1):
    spread = (asks1[0][0] - bids1[0][0]) / bids1[0][0]
    pos_diff, agg_pos, working_qty, price_edge = \
        self._get_target_position_info(implied_rate, spread)
    if abs(pos_diff) < self._lot_size * 1e-2:
      return

    # Buy only when ref price move up.
    # Sell only when ref price move down.
    move_threshold = ref_ma * self._trading_threshold
    ref_price_up = mid0 >= ref_ma + move_threshold
    ref_price_down = mid0 <= ref_ma - move_threshold

    order_qty = self._tick_info1[1].floor(abs(pos_diff))
    aggressive_amt = self._aggressive_amt
    if pos_diff > 0:
      if not ref_price_up:
        # self.logger.info(
        #    '%s ref price not up %f %f' % (self._ticker1, mid0, ref_ma))
        return
      order_side = OrderSide.BUY
      order_price = self._get_agg_buy_price(asks1[0][0], aggressive_amt)
      if asks1[0][0] + price_edge >= asks1[1][0]:
        order_price = max(order_price, asks1[1][0])
      order_qty *= (1 + self._agg_taker_fee)
    else:
      if not ref_price_down:
        # self.logger.info(
        #    '%s ref price not down %f %f' % (self._ticker1, mid0, ref_ma))
        return
      order_side = OrderSide.SELL
      order_price = self._get_agg_sell_price(bids1[0][0], aggressive_amt)
      if bids1[0][0] - price_edge <= bids1[1][0]:
        order_price = min(order_price, bids1[1][0])

    if not FLAGS.dry_run:
      ret = self._agg_executor1.submit_agg_order(order_side=order_side,
                                                 order_price=order_price,
                                                 order_qty=order_qty)
      if ret:
        if order_side == OrderSide.BUY:
          self._expected_working_qty += order_qty
        else:
          self._expected_working_qty -= order_qty
        self._last_agg_submit_fill_ts = self._get_ts()
      self.logger.info('agg submitted return code %s' % ret)
      self._fillinfo_writer.write(
          FillInfo('agg submitted',
                   symbol=self._product1.symbol,
                   exchange=self._product1.exchange,
                   fill_price=order_price,
                   fill_qty=order_qty,
                   order_side=order_side,
                   timestamp=self._get_ts()))
    self.logger.info('%s target pos %f, current pos %f, working qty %f, expected working qty %f' %
                     (self._ticker1,
                      pos_diff + agg_pos + working_qty,
                      agg_pos,
                      working_qty,
                      self._expected_working_qty))
    self.logger.info('%s %s %s %s' % (bids0[0], asks0[0], bids1[0], asks1[0]))
    self.logger.info('mean %f, std %f, current mean %f, size %d' %
                     (self._mean, self._std, self._ma.get_latest_data(), self._ma.size()))

  def on_book_record_only(self, exchange_id, book):
    self._book_ts[exchange_id] = book.timestamp
    if not self._is_book_valid():
      self.logger.warning('%s book is not updated %s %s.' %
                          (self._ticker0, self._book_ts[0], self._book_ts[1]))
      return
    bids0, asks0 = self.feed[self._product0]
    bids1, asks1 = self.feed[self._product1]
    book_changed = self._book_changed(bids0, asks0, self._last_book0) \
        if exchange_id == 0 else self._book_changed(bids1, asks1, self._last_book1)
    if book_changed:
      mid0 = self._parent.get_mid(bids0, asks0)
      mid1 = self._parent.get_mid(bids1, asks1)
      implied_rate = mid1 / mid0
      self._update_deque(implied_rate, mid0, mid1)
    if exchange_id == 0:
      self._last_book0 = (bids0, asks0)
    else:
      self._last_book1 = (bids1, asks1)

  def output_deque(self):
    file_name = '%s_%s_%s' % \
        (self._product0.exchange, self._product1.exchange, self._ticker0)
    self.logger.info('Start writing deque %s' % file_name)
    with open(file_name, 'w') as fo:
      content = '\n'.join(['%d,%f' % (e[0], e[1]) for e in self._ma._dq])
      fo.write(content)
    self.logger.info('Finish writing deque')

  def _book_changed(self, bids, asks, last_book):
    if last_book:
      last_bids, last_asks = last_book
      for i in range(self._book_levels):
        if np.any(bids[i] != last_bids[i]) or np.any(asks[i] != last_asks[i]):
          return True
    return False

  def on_book(self, exchange_id, book):
    self._book_ts[exchange_id] = book.timestamp
    if self._agg_executor1:
      if not self._is_book_valid():
        self.logger.warning('%s book is not updated %s %s.' %
                            (self._ticker0, self._book_ts[0], self._book_ts[1]))
        return

      bids0, asks0 = self.feed[self._product0]
      bids1, asks1 = self.feed[self._product1]
      book_changed = self._book_changed(bids0, asks0, self._last_book0) \
          if exchange_id == 0 else self._book_changed(bids1, asks1, self._last_book1)
      if book_changed:
        mid0 = self._parent.get_mid(bids0, asks0)
        mid1 = self._parent.get_mid(bids1, asks1)
        implied_rate = mid1 / mid0
        ref_ma = self._ref_ma
        self._update_deque(implied_rate, mid0, mid1)

        if self.ready():
          self._process(implied_rate, ref_ma, bids0, asks0, bids1, asks1, mid0, mid1)
      if not FLAGS.dry_run:
        self._agg_executor1.manage_agg_orders()
      if exchange_id == 0:
        self._last_book0 = (bids0, asks0)
      else:
        self._last_book1 = (bids1, asks1)

  def on_aggressive_fill(self, fill_price, fill_qty, order):
    self._agg_working_order_fills[order.order_id] = \
        self._agg_working_order_fills.get(order.order_id, 0.0) + fill_qty
    self._fillinfo_writer.write(
        FillInfo('agg filled',
                 symbol=order.product.symbol,
                 exchange=order.product.exchange,
                 fill_price=fill_price,
                 fill_qty=fill_qty,
                 order_side=order.order_side,
                 timestamp=self._get_ts()))
    if self._is_buy_order_func(order.order_side):
      self._update_pos1(self._pos1 + fill_qty * (1 - self._agg_taker_fee))
      self._expected_working_qty -= fill_qty
    else:
      self._update_pos1(self._pos1 - fill_qty)
      self._expected_working_qty += fill_qty
    self._last_agg_submit_fill_ts = self._get_ts()
    self.logger.info('expected %s pos, %f' % (self._ticker0, self._pos1))

  def on_trade(self, exchange_id, trade):
    pass


class BasisUnhedgedStrategy(BaseStrategy):
  def __init__(self, strategy_params, trading_tickers, record_implied_rate_only=False):
    self.logger = logging.getLogger(__name__)

    self._trading_tickers = set(trading_tickers)
    self.logger.info('Traded tickers %s' % self._trading_tickers)
    self._record_implied_rate_only = record_implied_rate_only
    self.logger.info('Record implied rate only %s' % self._record_implied_rate_only)

    self._config = strategy_params['config']

    self._exch0_subreq = SubscriptionRequest.from_dict(
        BaseStrategy.update_sub_config(self._config['common']['exch0_subreq']))
    self._exch1_subreq = SubscriptionRequest.from_dict(
        BaseStrategy.update_sub_config(self._config['common']['exch1_subreq']))

    self._exchange0 = self._exch0_subreq.exchange
    self._exchange1 = self._exch1_subreq.exchange
    self._og_config = strategy_params['og_config']

    self._balance_checker = None
    account_checker_config = self._get_account_checker_config()
    if account_checker_config is not None:
      self._balance_checker = AccountChecker(account_checker_config)

    self._book_builder_clses = (get_flow_book_builder_cls(self._exch0_subreq),
                                get_flow_book_builder_cls(self._exch1_subreq))

    self._live_env_clses1 = get_live_env_cls(self._exch1_subreq)
    self._live_env1 = None
    self.og1 = None

    self._products0, self._products1 = _gen_arb_products(self._config['common'],
                                                         self._config['products0'],
                                                         self._config['products1'],
                                                         self._book_builder_clses[0].ProductType,
                                                         self._book_builder_clses[1].ProductType)
    self._is_future = self._products1.products[0].product_type == 'Futures'

    exch0_exchange_feed_checker = self._config['common'].get('exch0_exchange_feed_checker', False)
    exch1_exchange_feed_checker = self._config['common'].get('exch1_exchange_feed_checker', False)
    feed_checker_products = []
    if exch0_exchange_feed_checker:
      feed_checker_products += self._products0.products
    if exch1_exchange_feed_checker:
      feed_checker_products += self._products1.products
    feed_checker_config = get_default_feed_checker_config(feed_checker_products)

    multi_feed_subsystem_cls = BaseStrategy.get_multi_feed_subsystem()
    self._feed = multi_feed_subsystem_cls(feed_checker_config=feed_checker_config)
    self._ts = None
    self._order_canceled_after_og_ready = False
    self._exit_attempt = 0
    self._exit_executor = None
    self._balance_printer = None
    self._fillinfo_writer = FillInfoWriter('fills.csv', FLAGS.strategy_name)
    self._fill_dumper = FillLogDumper(self._exch1_subreq)

    self._book_levels = FLAGS.book_levels
    self._last_feed = {}
    self._book_map = BookMap(self._products0.products + self._products1.products,
                             logger=self.logger)

    self._reserve_map = self._get_reserve_map()
    self._strategy_request = convert_strategy_request_from_dict(
        self._config['common']['strategy_request'])
    self._strategy_logger = StrategyLogProvider(self._strategy_request)
    self._reserve_proto = reserve_as_proto(self._exch1_subreq, self._reserve_map)

    self._substrategies = {}
    self._usubstrategies = []
    for product0, product1 in zip(self._products0.products, self._products1.products):
      strat = BasisUnhedgedSubStrategy(self,
                                       product0,
                                       product1,
                                       self._fillinfo_writer,
                                       self._strategy_logger)
      self._substrategies[product0] = self._substrategies[product1] = strat
      self._usubstrategies.append(strat)

  def periodic_write_logs(self):
    self.og1.og_logger.write_account_request()
    self._strategy_logger.write_request()
    self._strategy_logger.write_reserve(self._reserve_proto)
    for strat in self._usubstrategies:
      strat.write_executor_config()
    self._ioloop.add_timeout(datetime.timedelta(hours=1), self.periodic_write_logs)

  @property
  def strategy_logger(self):
    return self._strategy_logger

  @staticmethod
  def get_mid(bids, asks):
    return (bids[0][0] + asks[0][0]) / 2.0

  def _update_prices(self, product, book):
    bids = np.array(book.get_bid_array(self._book_levels))
    asks = np.array(book.get_ask_array(self._book_levels))
    self._last_feed[product] = (bids, asks)

  def get_implied_rate(self):
    rates = [strat.get_implied_rate() for strat in self._substrategies.values()]
    return 1.0 / np.nanmean(rates)

  def on_book(self, product, exchange_id, book):
    self._ts = book.timestamp
    self._update_prices(product, book)
    self._substrategies[product].on_book(exchange_id, book)
    self._book_map.set_book(product, book)

  def on_book_record_only(self, product, exchange_id, book):
    self._ts = book.timestamp
    self._update_prices(book)
    self._substrategies[product].on_book_record_only(exchange_id, book)

  def on_trade(self, exchange_id, trade):
    pass

  def on_aggressive_fill(self, fill_price, fill_qty, order):
    self._substrategies[order.product].on_aggressive_fill(fill_price, fill_qty, order)
    self._fill_dumper.dump_fill(fill_price, fill_qty, order)

  def _get_account_checker_config(self):
    if self._config.get('account_checker') is None:
      return None

    config_keys = {
        'reserve',
        'lot_size',
        'stack_size',
        'ratio_min_balance',
        'ratio_max_balance',
        'ratio_min_hold_balance',
    }

    balance_check_config_list = []
    for currency, config in self._config['account_checker'].items():
      config = {**config}
      if 'lot_size' not in config:
        for product in self._config['products']:
          base = SharedProduct.FromStrNativeProduct(product).base.currency
          if base == currency:
            config['lot_size'] = self._config['products'][product]['lot_size']
            break

      assert config_keys == set(config.keys()), config
      reserve = config['reserve']
      config['currency'] = currency
      config['min_balance'] = reserve * config['ratio_min_balance']
      config['max_balance'] = reserve * config['ratio_max_balance']
      config['min_available_balance'] = 0.0
      config['max_available_balance'] = config['lot_size'] * config['stack_size']
      config['min_hold_balance'] = reserve * config['ratio_min_hold_balance']
      config['max_hold_balance'] = reserve * config['ratio_max_balance']
      currency_checker_config = convert_checker_config_from_dict(config)
      balance_check_config_list.append(currency_checker_config)

    return AccountCheckerConfig(each_balance_check_config=balance_check_config_list)

  def signal_handler(self, sig, frame):
    self.logger.warning('Caught signal: %s', sig)
    self._ioloop.add_callback_from_signal(self.on_exit)

  def on_exit(self):
    self._exit_attempt += 1
    if self._exit_attempt > 5:
      self._ioloop.add_callback(self._ioloop.stop)

    finished = True
    if self._exit_executor is not None and not FLAGS.dry_run:
      ogs = to_list([self.og1])
      self.logger.info('Exit callback executed.')
      finished = self._exit_executor.cancel_all_on_exit(ogs)

    if finished:
      self._ioloop.add_callback(self._ioloop.stop)
    else:
      timeout = datetime.timedelta(seconds=2)
      self._ioloop.add_timeout(timeout, self.on_exit)

  def start(self, ioloop, queue, flow_sub):
    self._ioloop = ioloop
    self._feed.start(ioloop, queue, flow_sub)
    if not self._record_implied_rate_only:
      self._live_env1.start()
      self.start_print_states()
    else:
      self.output_deque(periodic=True)

  def subscribe_feeds(self):
    if self._record_implied_rate_only:
      func = self.on_book_record_only
    else:
      func = self.on_book
    for product in self._products0.products:
      self._feed.subscribe_book(FeedSubscriptionRequest.create([product], self._exch0_subreq),
                                functools.partial(func, product, 0))
    for product in self._products1.products:
      self._feed.subscribe_book(FeedSubscriptionRequest.create([product], self._exch1_subreq),
                                functools.partial(func, product, 1))
    if FLAGS.subscribe_trade:
      self._feed.subscribe_trade(self._products0.products, functools.partial(self.on_trade, 0))
      self._feed.subscribe_trade(self._products1.products, functools.partial(self.on_trade, 1))

  def subscribe_live_envs(self, ioloop, queue, flow_sub):
    og_params = self._config.get('og_params', {})
    og_params1 = convert_og_config_from_dict(self._exch1_subreq, og_params)

    self._live_env1 = create_live_env_instance(self._live_env_clses1,
                                               self._products1.currencies,
                                               self._products1.products,
                                               self._og_config[0],
                                               og_params=og_params1,
                                               ioloop=ioloop,
                                               queue=queue,
                                               flow_sub=flow_sub,
                                               og_log_callback=self._strategy_logger.write_og_log)
    self._live_env1.on_order_gateway_reset_callback.append(self.register_order_gateway)

  def _get_reserve_map(self):
    reserves = dict()
    for product in self._products1.products:
      key = BasisUnhedgedStrategy._get_config_key(product)
      reserves[product.base.currency] = (self._config['products1'][key]['max_pos']
                                         + self._config['products1'][key]['min_pos']) / 2.0
    return reserves

  def register_balance_printer(self):
    printer_cls = get_pnl_balance_printer_cls(self._exch1_subreq)
    self._balance_printer = printer_cls(self._products1.products,
                                        self._book_map,
                                        self._reserve_map,
                                        logger=self.logger)
    self._balance_printer.register_order_gateway(self.og1)
    PeriodicCallback(self.dump_pnl_balance, 30 * 1000).start()

  def start_print_states(self):
    for strat in self._usubstrategies:
      states = strat.get_states()
      self.logger.info(
          '%s mean: %s, std: %s, current_mean: %s, size: %d' %
          (strat._ticker1, states['mean'], states['std'], states['current_mean'],
           states['dq_size']))
    self._ioloop.add_timeout(datetime.timedelta(seconds=10), self.start_print_states)

  def dump_pnl_balance(self, periodic=False):
    if self._balance_printer is not None:
      self._balance_printer.print_pnl_balance()
      pnl_balances = self._balance_printer.calculate_pnl_balance()
      if pnl_balances is not None:
        for pnl_balance in pnl_balances:
          self._strategy_logger.write_pnl_balance(pnl_balance=pnl_balance,
                                                  quote=None,
                                                  subreq_id=None)

  def output_deque(self, periodic=False):
    for strat in set(self._substrategies.values()):
      strat.output_deque()
    if periodic:
      self._ioloop.add_timeout(datetime.timedelta(seconds=30), self.output_deque, True)

  @staticmethod
  def _get_config_key(product):
    key = product.subscription_symbol
    if product.exchange == 'Huobi' and product.product_type == 'Futures':
      key = key.replace('CURRENT', 'THIS')
      key = key.replace('THIS_QUARTER', 'QUARTER')
    return key

  def register_order_gateway(self, name, order_gateway):
    self.logger.info('OG initialized: %s', name)
    if order_gateway.exchange == self._exchange1:
      self.og1 = order_gateway
    # cancel working orders first after OGs are ready.
    if not self._order_canceled_after_og_ready:
      order_gateway.cancel_all()
      self.logger.info('Orders canceled after OGs are ready.')
      self._order_canceled_after_og_ready = True
    self._exit_executor = SimpleCancelAllExecutor()

    if self.og1:
      self.register_balance_printer()
      self.og1.register_fill_callback(self.on_aggressive_fill)
      for prod0, prod1 in zip(self._products0.products, self._products1.products):
        key = BasisUnhedgedStrategy._get_config_key(prod1)
        config1 = self._config['products1'][key].copy()
        config1.update(self._config.get('agg_executor_configs', {}))
        self._substrategies[prod1].init_executors(self._exch1_subreq, config1, self.og1)
      self.periodic_write_logs()


def get_param_set(og_configs, config_filename):
  og_config = og_configs.split(',')
  config = load_clean_json(config_filename)
  return {
      'og_config': og_config,
      'config': config,
  }


def main(argv):
  assert FLAGS.strategy_name

  strategy_params = get_param_set(og_configs=FLAGS.og_configs,
                                  config_filename=FLAGS.config_filename)

  queue = SimpleQueue('basisunhedged')
  flow_sub = flow_subscriber.from_queue(queue)
  ioloop = IOLoop.current()

  record_implied_rate_only = FLAGS.record_implied_rate_only

  trading_tickers = FLAGS.trading_tickers.split(',')
  strategy = BasisUnhedgedStrategy(strategy_params,
                                   trading_tickers,
                                   record_implied_rate_only=record_implied_rate_only)
  strategy.subscribe_feeds()
  if not record_implied_rate_only:
    strategy.subscribe_live_envs(ioloop, queue, flow_sub)
  strategy.start(ioloop, queue, flow_sub)

  # exit_after = datetime.timedelta(days=1)
  # ioloop.add_timeout(exit_after, strategy.on_exit)

  signal.signal(signal.SIGTERM, strategy.signal_handler)
  signal.signal(signal.SIGINT, strategy.signal_handler)

  ioloop.start()


if __name__ == '__main__':
  flags.DEFINE_string('og_configs', '', '')

  flags.DEFINE_string('config_filename', '', '')

  flags.DEFINE_integer('exit_after_min', 30, '')

  flags.DEFINE_boolean('dry_run', False, '')

  flags.DEFINE_string('trading_tickers', '', '')

  flags.DEFINE_float('aggressive_bp', 3, '')

  flags.DEFINE_float('book_invalid_sec', 3.5, '')

  flags.DEFINE_boolean('update_agg_pos_with_balance', False, '')

  flags.DEFINE_float('update_pos_with_balance_sec', 4, '')

  flags.DEFINE_float('update_working_qty_sec', 3, '')

  flags.DEFINE_float('window_period_sec', 600, '')

  flags.DEFINE_float('ref_price_window_period_sec', 5, '')

  flags.DEFINE_float('enter_position_threshold', 3.0, '')

  flags.DEFINE_float('exit_position_threshold', -1.5, '')

  flags.DEFINE_float('spread_multiplier', 0.5, '')

  flags.DEFINE_boolean('record_implied_rate_only', False, '')

  flags.DEFINE_float('trading_threshold_bp', 20, '')

  flags.DEFINE_integer('book_levels', 3, '')

  flags.DEFINE_float('max_std_ratio', 10e-4, '')

  flags.DEFINE_float('min_std_ratio', 5e-4, '')

  flags.DEFINE_string('strategy_name', '', '')

  flags.DEFINE_boolean('subscribe_trade', False, '')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  app.run(main)
